{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Find the Kth Smallest Sum of a Matrix With Sorted Rows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #array #binary-search #matrix #heap-priority-queue"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #数组 #二分查找 #矩阵 #堆（优先队列）"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: kthSmallest"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #有序矩阵中的第 k 个最小数组和"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给你一个 <code>m&nbsp;* n</code> 的矩阵 <code>mat</code>，以及一个整数 <code>k</code> ，矩阵中的每一行都以非递减的顺序排列。</p>\n",
    "\n",
    "<p>你可以从每一行中选出 1 个元素形成一个数组。返回所有可能数组中的第 k 个 <strong>最小</strong> 数组和。</p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>示例 1：</strong></p>\n",
    "\n",
    "<pre><strong>输入：</strong>mat = [[1,3,11],[2,4,6]], k = 5\n",
    "<strong>输出：</strong>7\n",
    "<strong>解释：</strong>从每一行中选出一个元素，前 k 个和最小的数组分别是：\n",
    "[1,2], [1,4], [3,2], [3,4], [1,6]。其中第 5 个的和是 7 。  </pre>\n",
    "\n",
    "<p><strong>示例 2：</strong></p>\n",
    "\n",
    "<pre><strong>输入：</strong>mat = [[1,3,11],[2,4,6]], k = 9\n",
    "<strong>输出：</strong>17\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 3：</strong></p>\n",
    "\n",
    "<pre><strong>输入：</strong>mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7\n",
    "<strong>输出：</strong>9\n",
    "<strong>解释：</strong>从每一行中选出一个元素，前 k 个和最小的数组分别是：\n",
    "[1,1,2], [1,1,3], [1,4,2], [1,4,3], [1,1,6], [1,5,2], [1,5,3]。其中第 7 个的和是 9 。 \n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 4：</strong></p>\n",
    "\n",
    "<pre><strong>输入：</strong>mat = [[1,1,10],[2,2,9]], k = 7\n",
    "<strong>输出：</strong>12\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>m == mat.length</code></li>\n",
    "\t<li><code>n == mat.length[i]</code></li>\n",
    "\t<li><code>1 &lt;= m, n &lt;= 40</code></li>\n",
    "\t<li><code>1 &lt;= k &lt;= min(200, n ^&nbsp;m)</code></li>\n",
    "\t<li><code>1 &lt;= mat[i][j] &lt;= 5000</code></li>\n",
    "\t<li><code>mat[i]</code> 是一个非递减数组</li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows](https://leetcode.cn/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows](https://leetcode.cn/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[[1,3,11],[2,4,6]]\\n5', '[[1,3,11],[2,4,6]]\\n9', '[[1,10,10],[1,4,5],[2,3,6]]\\n7']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "class Solution:\n",
    "\n",
    "    def kSmallesPairs(self, num1, num2, k):\n",
    "        ans = []\n",
    "        h = [(num1[0] + num2[0], 0, 0)]\n",
    "        while h and len(ans) < k:\n",
    "            _, i, j = heappop(h)\n",
    "            ans.append(num1[i] + num2[j])\n",
    "            if j == 0 and i + 1 < len(num1):\n",
    "                heappush(h, (num1[i + 1] + num2[0], i + 1, 0))\n",
    "            if j + 1 < len(num2):\n",
    "                heappush(h, (num1[i] + num2[j + 1], i, j + 1))\n",
    "        return ans\n",
    "\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        a = mat[0][:k]\n",
    "        for row in mat[1:]:\n",
    "            a = self.kSmallesPairs(row, a, k)\n",
    "\n",
    "        return a[-1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import heapq\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        pre = mat[0]\n",
    "        for i in range(1, len(mat)):\n",
    "            pre = self.count(pre, mat[i], k)\n",
    "            #return\n",
    "        return pre[k-1]\n",
    "    \n",
    "    def count(self, x, y, k):\n",
    "        if len(x) < len(y):\n",
    "            return self.count(y, x, k)\n",
    "        #h1 x[i] + y[0]  h2 x[i] + y[1]\n",
    "        h = [(x[0] + y[i], 0) for i in range(len(y))]\n",
    "        heapq.heapify(h)\n",
    "        res = []\n",
    "        while len(res) < k and len(h) > 0:\n",
    "            num, idx = heapq.heappop(h)\n",
    "            #print(num, idx, h)\n",
    "            res.append(num)\n",
    "            if idx + 1 < len(x):\n",
    "                heapq.heappush(h, (x[idx + 1] - x[idx] + num, idx + 1))\n",
    "        return res\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    # 373. 查找和最小的 K 对数字\n",
    "    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[int]:\n",
    "        ans = []\n",
    "        h = [(nums1[0] + nums2[0], 0, 0)]\n",
    "        while h and len(ans) < k:\n",
    "            _, i, j = heappop(h)\n",
    "            ans.append(nums1[i] + nums2[j])  # 数对和\n",
    "            if j == 0 and i + 1 < len(nums1):\n",
    "                heappush(h, (nums1[i + 1] + nums2[0], i + 1, 0))\n",
    "            if j + 1 < len(nums2):\n",
    "                heappush(h, (nums1[i] + nums2[j + 1], i, j + 1))\n",
    "        return ans\n",
    "\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        a = mat[0][:k]\n",
    "        for row in mat[1:]:\n",
    "            a = self.kSmallestPairs(row, a, k)\n",
    "        return a[-1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import heapq\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        def merge(f: List[int], g: List[int], k: int) -> List[int]:\n",
    "            if len(g) > len(f):\n",
    "                return merge(g, f, k)\n",
    "            \n",
    "            q = [(f[0] + g[i], 0, i) for i in range(len(g))]\n",
    "            heapq.heapify(q)\n",
    "\n",
    "            ans = list()\n",
    "            while k and q:\n",
    "                entry = heapq.heappop(q)\n",
    "                ans.append(entry[0])\n",
    "                if entry[1] + 1 < len(f):\n",
    "                    heapq.heappush(q, (f[entry[1] + 1] + g[entry[2]], entry[1] + 1, entry[2]))\n",
    "                k -= 1\n",
    "            \n",
    "            return ans\n",
    "        \n",
    "        prev = mat[0]\n",
    "        for i in range(1, len(mat)):\n",
    "            prev = merge(prev, mat[i], k)\n",
    "        return prev[k - 1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        def merge(f: List[int], g: List[int], k: int) -> List[int]:\n",
    "            if len(g) > len(f):\n",
    "                return merge(g, f, k)\n",
    "            \n",
    "            q = [(f[0] + g[i], 0, i) for i in range(len(g))]\n",
    "            heapq.heapify(q)\n",
    "\n",
    "            ans = list()\n",
    "            while k and q:\n",
    "                entry = heapq.heappop(q)\n",
    "                ans.append(entry[0])\n",
    "                if entry[1] + 1 < len(f):\n",
    "                    heapq.heappush(q, (f[entry[1] + 1] + g[entry[2]], entry[1] + 1, entry[2]))\n",
    "                k -= 1\n",
    "            \n",
    "            return ans\n",
    "        \n",
    "        prev = mat[0]\n",
    "        for i in range(1, len(mat)):\n",
    "            prev = merge(prev, mat[i], k)\n",
    "        return prev[k - 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    # 373. 查找和最小的 K 对数字\n",
    "    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[int]:\n",
    "        ans = []\n",
    "        h = [(nums1[0] + nums2[0], 0, 0)]\n",
    "        while h and len(ans) < k:\n",
    "            _, i, j = heappop(h)\n",
    "            ans.append(nums1[i] + nums2[j])  # 数对和\n",
    "            if j == 0 and i + 1 < len(nums1):\n",
    "                heappush(h, (nums1[i + 1] + nums2[0], i + 1, 0))\n",
    "            if j + 1 < len(nums2):\n",
    "                heappush(h, (nums1[i] + nums2[j + 1], i, j + 1))\n",
    "        return ans\n",
    "\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        a = mat[0][:k]\n",
    "        for row in mat[1:]:\n",
    "            a = self.kSmallestPairs(row, a, k)\n",
    "        return a[-1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        pre = [0]\n",
    "        for cur in mat:\n",
    "            pre = sorted(a + b for a in pre\n",
    "                                for b in cur[:k])[:k]\n",
    "        return pre[-1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from queue import *\n",
    "class Solution:\n",
    "\n",
    "    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:\n",
    "        pq = PriorityQueue()\n",
    "        m,n = len(nums1),len(nums2)\n",
    "        for i in range(min(m,k)):\n",
    "            pq.put((nums1[i] + nums2[0], i, 0))\n",
    "        ans = []\n",
    "        while k and not pq.empty():\n",
    "            k -= 1\n",
    "            _,i, j = pq.get()\n",
    "            ans.append(nums1[i] + nums2[j])\n",
    "            if j + 1 < n:\n",
    "                pq.put((nums1[i] + nums2[j + 1],i, j+1))\n",
    "        return ans\n",
    "\n",
    "\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        a = mat[0][:k]\n",
    "        for row in mat[1:]:\n",
    "            a = self.kSmallestPairs(row, a, k)\n",
    "        return a[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import heapq\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        pre = mat[0]\n",
    "        for i in range(1, len(mat)):\n",
    "            pre = self.count(pre, mat[i], k)\n",
    "            #return\n",
    "        return pre[k-1]\n",
    "    \n",
    "    def count(self, x, y, k):\n",
    "        if len(x) > len(y):\n",
    "            return self.count(y, x, k)\n",
    "        #h1 x[i] + y[0]  h2 x[i] + y[1]\n",
    "        h = [(x[0] + y[i], 0) for i in range(len(y))]\n",
    "        heapq.heapify(h)\n",
    "        res = []\n",
    "        while len(res) < k and len(h) > 0:\n",
    "            num, idx = heapq.heappop(h)\n",
    "            #print(num, idx, h)\n",
    "            res.append(num)\n",
    "            if idx + 1 < len(x):\n",
    "                heapq.heappush(h, (x[idx + 1] - x[idx] + num, idx + 1))\n",
    "        return res\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        def merge(a, b, k):\n",
    "            h = [(a[i]+b[0], i, 0) for i in range(min(len(a), k))]\n",
    "            heapify(h)\n",
    "            res = []\n",
    "            while h:\n",
    "                _, i, j = heappop(h)\n",
    "                res.append(a[i]+b[j])\n",
    "                if len(res) == k:\n",
    "                    return res\n",
    "                if j+1 < len(b):\n",
    "                    heappush(h, (a[i]+b[j+1], i, j+1))\n",
    "            return res\n",
    "        \n",
    "        prev = mat[0]\n",
    "        for i in range(1, len(mat)):\n",
    "            prev = merge(prev, mat[i], k)\n",
    "        \n",
    "        return prev[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def __init__(self):\n",
    "        self.mat, self.cnt = [], 1\n",
    "        self.m, self.n = 1, 1\n",
    "        self.k = 1\n",
    "\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        self.mat = mat\n",
    "        self.m, self.n, self.k = len(mat), len(mat[0]), k\n",
    "\n",
    "        left, right = 0, 0\n",
    "        for i in range(self.m):\n",
    "            left += mat[i][0]\n",
    "            right += mat[i][-1]\n",
    "\n",
    "        init_sum = left\n",
    "        while left < right:\n",
    "            mid = left + (right - left) // 2\n",
    "            self.cnt = 1\n",
    "            self.dfs(0, init_sum, mid)\n",
    "            if self.cnt < k:\n",
    "                left = mid + 1\n",
    "            else:\n",
    "                right = mid\n",
    "        return left\n",
    "\n",
    "    def dfs(self, i, cur_sum, target_sum):\n",
    "        if i == self.m or self.cnt >= self.k:\n",
    "            return\n",
    "\n",
    "        for j in range(self.n):\n",
    "            new_sum = cur_sum + self.mat[i][j] - self.mat[i][0]\n",
    "            if new_sum > target_sum:\n",
    "                break\n",
    "            if j > 0:\n",
    "                self.cnt += 1\n",
    "            self.dfs(i + 1, new_sum, target_sum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def __init__(self):\n",
    "        self.mat, self.cnt = [], 1\n",
    "        self.m, self.n = 1, 1\n",
    "        self.k = 1\n",
    "\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        self.mat = mat\n",
    "        self.m, self.n, self.k = len(mat), len(mat[0]), k\n",
    "\n",
    "        left, right = 0, 0\n",
    "        for i in range(self.m):\n",
    "            left += mat[i][0]\n",
    "            right += mat[i][-1]\n",
    "\n",
    "        init_sum = left\n",
    "        while left < right:\n",
    "            mid = left + (right - left) // 2\n",
    "            self.cnt = 1\n",
    "            self.dfs(0, init_sum, mid)\n",
    "            if self.cnt < k:\n",
    "                left = mid + 1\n",
    "            else:\n",
    "                right = mid\n",
    "        return left\n",
    "\n",
    "    def dfs(self, row: int, cur_sum: int, target: int):\n",
    "        if row == self.m or self.cnt >= self.k: return\n",
    "\n",
    "        for j in range(self.n):\n",
    "            new_sum = cur_sum + self.mat[row][j] - self.mat[row][0]\n",
    "            if new_sum <= target:\n",
    "                if j > 0: self.cnt += 1\n",
    "                self.dfs(row + 1, new_sum, target)\n",
    "            else:\n",
    "                break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def __init__(self):\n",
    "        self.mat, self.cnt = [], 1\n",
    "        self.m, self.n = 1, 1\n",
    "        self.k = 1\n",
    "\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        self.mat = mat\n",
    "        self.m, self.n, self.k = len(mat), len(mat[0]), k\n",
    "\n",
    "        left, right = 0, 0\n",
    "        for i in range(self.m):\n",
    "            left += mat[i][0]\n",
    "            right += mat[i][-1]\n",
    "\n",
    "        init_sum = left\n",
    "        while left < right:\n",
    "            mid = left + (right - left) // 2\n",
    "            self.cnt = 1\n",
    "            self.dfs(0, init_sum, mid)\n",
    "            if self.cnt < k:\n",
    "                left = mid + 1\n",
    "            else:\n",
    "                right = mid\n",
    "        return left\n",
    "\n",
    "    def dfs(self, row: int, cur_sum: int, target: int):\n",
    "        if row == self.m or self.cnt >= self.k: return\n",
    "\n",
    "        for j in range(self.n):\n",
    "            new_sum = cur_sum + self.mat[row][j] - self.mat[row][0]\n",
    "            if new_sum <= target:\n",
    "                if j > 0: self.cnt += 1\n",
    "                self.dfs(row + 1, new_sum, target)\n",
    "            else:\n",
    "                break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "class Solution:\n",
    "\n",
    "    def kSmallesPairs(self, num1, num2, k):\n",
    "        ans = []\n",
    "        # num1 是 下一行， num2 是之前行的最小值的组合的和\n",
    "        # 都是递增的，我们管理下一行的答案\n",
    "        h = [(num1[0] + num2[0], 0, 0)]  # 小根堆 # 插入小根堆，第一个为和，第二第三为在两个数组的和\n",
    "        while h and len(ans) < k:\n",
    "            _, i, j = heappop(h)  # 最小的数\n",
    "            # 只有堆中最小的数才能进行添加入 答案 ans\n",
    "            ans.append(num1[i] + num2[j])  # 加入该数\n",
    "            if j == 0 and i + 1 < len(num1):\n",
    "                # 继续往堆里面加数，这时候的堆中可以添加两个数，一个是 目前还是第一个 j\n",
    "                # 当最小值 还是 第一个数的时候，我们这时候需要往num1继续取值\n",
    "                heappush(h, (num1[i + 1] + num2[0], i + 1, 0))\n",
    "            if j + 1 < len(num2):\n",
    "                # 如果不是的话，我们往 j + 1 取值，往答案下一个取值，放入堆，堆中元素进行管理\n",
    "                heappush(h, (num1[i] + num2[j + 1], i, j + 1))\n",
    "        return ans\n",
    "\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        a = mat[0][:k]\n",
    "        for row in mat[1:]:\n",
    "            a = self.kSmallesPairs(row, a, k)\n",
    "\n",
    "        return a[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        a = mat[0][:k]\n",
    "        for row in mat[1:]:\n",
    "            a = sorted(x + y for x in a for y in row)[:k]\n",
    "        return a[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[int]:\n",
    "        ans = []\n",
    "        h = [(nums1[0] + nums2[0], 0, 0)]\n",
    "        while h and len(ans) < k:\n",
    "            _, i, j = heappop(h)\n",
    "            ans.append(nums1[i] + nums2[j])  # 数对和\n",
    "            if j == 0 and i + 1 < len(nums1):\n",
    "                heappush(h, (nums1[i + 1] + nums2[0], i + 1, 0))\n",
    "            if j + 1 < len(nums2):\n",
    "                heappush(h, (nums1[i] + nums2[j + 1], i, j + 1))\n",
    "        return ans\n",
    "\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        a = mat[0][:k]\n",
    "        for row in mat[1:]:\n",
    "            a = self.kSmallestPairs(row, a, k)\n",
    "        return a[-1]\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        def merge(f: List[int], g: List[int], k: int) -> List[int]:\n",
    "            left, right, thres = f[0] + g[0], f[-1] + g[-1], 0\n",
    "            k = min(k, len(f) * len(g))\n",
    "            while left <= right:\n",
    "                mid = (left + right) // 2\n",
    "                rptr, cnt = len(g) - 1, 0\n",
    "                for lptr, x in enumerate(f):\n",
    "                    while rptr >= 0 and x + g[rptr] > mid:\n",
    "                        rptr -= 1\n",
    "                    cnt += rptr + 1\n",
    "            \n",
    "                if cnt >= k:\n",
    "                    thres = mid;\n",
    "                    right = mid - 1;\n",
    "                else:\n",
    "                    left = mid + 1;\n",
    "\n",
    "            ans = list()\n",
    "            for i, fi in enumerate(f):\n",
    "                for j, gj in enumerate(g):\n",
    "                    if (total := fi + gj) < thres:\n",
    "                        ans.append(total)\n",
    "                    else:\n",
    "                        break\n",
    "            \n",
    "            ans += [thres] * (k - len(ans))\n",
    "            ans.sort()\n",
    "            return ans\n",
    "        \n",
    "        prev = mat[0]\n",
    "        for i in range(1, len(mat)):\n",
    "            prev = merge(prev, mat[i], k)\n",
    "        return prev[k - 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    # 373. 查找和最小的 K 对数字\n",
    "    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[int]:\n",
    "        ans = []\n",
    "        h = [(nums1[i] + nums2[0], i, 0) for i in range(len(nums1))]\n",
    "        while h and len(ans) < k:\n",
    "            _, i, j = heappop(h)\n",
    "            ans.append(nums1[i] + nums2[j])  # 数对和\n",
    "            if j + 1 < len(nums2):\n",
    "                heappush(h, (nums1[i] + nums2[j + 1], i, j + 1))\n",
    "        return ans\n",
    "\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        a = mat[0][:k]\n",
    "        for row in mat[1:]:\n",
    "            a = self.kSmallestPairs(row, a, k)\n",
    "        return a[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import heapq\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        pre = mat[0]\n",
    "        for i in range(1, len(mat)):\n",
    "            pre = self.count(pre, mat[i], k)\n",
    "            #return\n",
    "        return pre[k-1]\n",
    "    \n",
    "    def count(self, x, y, k):\n",
    "        #h1 x[i] + y[0]  h2 x[i] + y[1]\n",
    "        h = [(x[0] + y[i], 0) for i in range(len(y))]\n",
    "        heapq.heapify(h)\n",
    "        res = []\n",
    "        while len(res) < k and len(h) > 0:\n",
    "            num, idx = heapq.heappop(h)\n",
    "            print(num, idx, h)\n",
    "            res.append(num)\n",
    "            if idx + 1 < len(x):\n",
    "                heapq.heappush(h, (x[idx + 1] - x[idx] + num, idx + 1))\n",
    "        return res\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        a = mat[0][:k]\n",
    "        for row in mat[1:]:\n",
    "            a = sorted(x + y for x in a for y in row)[:k]\n",
    "        return a[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "      n = len(mat)\n",
    "      def check(s):\n",
    "        left_k = k\n",
    "        def dfs(i, s):\n",
    "          nonlocal left_k\n",
    "          if i < 0:\n",
    "            left_k -= 1\n",
    "            return left_k == 0\n",
    "          \n",
    "          for x in mat[i]:\n",
    "            if s - x + mat[i][0] < 0: break\n",
    "            if dfs(i - 1, s - x + mat[i][0]):\n",
    "              return True\n",
    "          return False\n",
    "        return dfs(n - 1, s - sl)\n",
    "      sl = sum(row[0] for row in mat)\n",
    "      sr = sum(row[-1] for row in mat)\n",
    "      left, right = sl, sr\n",
    "      while left <= right:\n",
    "        mid = (left + right)//2\n",
    "        if check(mid):\n",
    "          right = mid - 1\n",
    "        else:\n",
    "          left = mid + 1\n",
    "      return left\n",
    "          \n",
    "          \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        m , n = len(mat) , len(mat[0])\n",
    "        for i in range(m):\n",
    "            mat[i] = sorted(mat[i])\n",
    "        nums = []\n",
    "        for i in range(m):\n",
    "            if i == 0:\n",
    "                nums = copy.deepcopy(mat[i])\n",
    "            else:\n",
    "                tmp = []\n",
    "                for num in nums:\n",
    "                    for j in range(n):\n",
    "                        tmp.append(num+mat[i][j])\n",
    "                nums = tmp\n",
    "            nums = sorted(nums)\n",
    "            nums = nums[:k]\n",
    "        return nums[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import heapq\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        pre = mat[0]\n",
    "        for i in range(1, len(mat)):\n",
    "            pre = self.count(pre, mat[i], k)\n",
    "            #return\n",
    "        return pre[k-1]\n",
    "    \n",
    "    def count(self, x, y, k):\n",
    "        if len(x) < len(y):\n",
    "            return self.count(y, x, k)\n",
    "        #h1 x[i] + y[0]  h2 x[i] + y[1]\n",
    "        h = [(x[0] + y[i], 0) for i in range(len(y))]\n",
    "        heapq.heapify(h)\n",
    "        res = []\n",
    "        while len(res) < k and len(h) > 0:\n",
    "            num, idx = heapq.heappop(h)\n",
    "            print(num, idx, h)\n",
    "            res.append(num)\n",
    "            if idx + 1 < len(x):\n",
    "                heapq.heappush(h, (x[idx + 1] - x[idx] + num, idx + 1))\n",
    "        return res\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        m,n = len(mat), len(mat[0])\n",
    "        a = mat[0][:k]\n",
    "        for i in range(1,m):\n",
    "            a = sorted([x+y for x in a for y in mat[i]])[:k]\n",
    "        return a[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        pre = mat[0][0:k]\n",
    "        m = len(mat)\n",
    "        for i in range(1, m):\n",
    "            pre = sorted([prei + v for prei in pre for v in mat[i]])[0:k]\n",
    "        return pre[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        m = len(mat)\n",
    "        n = len(mat[0])\n",
    "        a = mat[0][:min(k,n)]#第0排前k个数字\n",
    "\n",
    "        for r  in range(1,m):\n",
    "            b = mat[r]\n",
    "            print(a,b)\n",
    "            a = self.kSmallestPairs(a,b,k)\n",
    "        return a[-1]\n",
    "\n",
    "    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:\n",
    "        #已经排序，在堆中更新nums2的index\n",
    "        m = len(nums1)\n",
    "        n = len(nums2)\n",
    "        nums= [(nums1[i]+nums2[0],i,0) for i in range(min(k,m))]#已经排序，无需heapify,nums1只需要前个数(k,m)\n",
    "        # print(nums)\n",
    "        # heapify(nums)\n",
    "        res = []\n",
    "        while k and nums:\n",
    "            he,i,j = heappop(nums)\n",
    "            res.append(he)\n",
    "            k -= 1\n",
    "            if j+1<n:\n",
    "                heappush(nums,(nums1[i]+nums2[j+1],i,j+1))\n",
    "            \n",
    "        return res "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        a=mat[0][:k]\n",
    "        n=len(mat)\n",
    "        for i in range(1,n):\n",
    "            arr=[]\n",
    "            for p in a:\n",
    "                for q in mat[i]:\n",
    "                    arr.append(p+q)\n",
    "            a=sorted(arr)[:k]\n",
    "        return a[k-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "\n",
    "\n",
    "# class Solution:\n",
    "#     def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "#         a = mat[0][:k]  # k '\n",
    "#         for row in mat[1:]: # 遍历后面的行\n",
    "#             a = sorted(x + y for x in a for y in row)[:k] # 每次都只取前面k个\n",
    "#         return a[-1]\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        def check(s):\n",
    "            left_k = k\n",
    "\n",
    "            def dfs(i, s):  # s 是剩余的个数\n",
    "                if i < 0:\n",
    "                    nonlocal left_k\n",
    "                    left_k -= 1\n",
    "                    return left_k == 0\n",
    "                for x in mat[i]:\n",
    "                    if x - mat[i][0] > s:\n",
    "                        break\n",
    "                    if dfs(i - 1, s - (x - mat[i][0])):  # 选择 mat[i] 中的每个数的值的结果\n",
    "                        return True\n",
    "                return False\n",
    "            return dfs(len(mat) - 1, s - sl)\n",
    "        sl = sum(row[0] for row in mat)\n",
    "        sr = sum(row[-1] for row in mat)\n",
    "        return sl + bisect_left(range(sl, sr), True, key=check)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        res = mat[0]\n",
    "        for line in mat[1:]: res = sorted([a+b for a in line for b in res])[:k]\n",
    "        return res[-1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        a = mat[0][:k]\n",
    "        for row in mat[1:]:\n",
    "            a = sorted(x + y for x in a for y in row[:k])[:k]\n",
    "        return a[k-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import copy\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        n = len(mat)\n",
    "        m = len(mat[0])\n",
    "        s0 = [sum([mat[i][0]for i in range(n)])]\n",
    "        stack = []\n",
    "        f = []\n",
    "        for i in range(n):\n",
    "            tmp = [0]*n\n",
    "            tmp[i] = 1\n",
    "            stack.append(tmp)\n",
    "            s0.append(sum([mat[j][tmp[j]] for j in range(n)]))\n",
    "            f.append(i)\n",
    "        s0.sort()\n",
    "        s0 = s0[:k]\n",
    "        while(len(stack)>0):\n",
    "            t = stack[0]\n",
    "            for i in range(f[0], n):\n",
    "                t_ = copy.copy(t)\n",
    "                t_[i] = t[i] + 1\n",
    "                if t_[i] < m:\n",
    "                    if sum([mat[j][t_[j]] for j in range(n)]) < s0[-1] or len(s0)<k:\n",
    "                        stack.append(t_)\n",
    "                        s0.append(sum([mat[j][t_[j]] for j in range(n)]))\n",
    "                        f.append(i)\n",
    "                        s0.sort()\n",
    "                        s0 = s0[:k]\n",
    "            stack = stack[1:]\n",
    "            f = f[1:]\n",
    "        return s0[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        def check(s: int) -> bool:\n",
    "            left_k = k\n",
    "            # 返回是否找到 k 个子数组和\n",
    "            def dfs(i: int, s: int) -> bool:\n",
    "                if i < 0:  # 能递归到这里，说明数组和不超过二分的 s\n",
    "                    nonlocal left_k\n",
    "                    left_k -= 1\n",
    "                    return left_k == 0  # 是否找到 k 个\n",
    "                for x in mat[i]:  # 「枚举选哪个」，注意 mat[i] 是有序的\n",
    "                    if x - mat[i][0] > s:  # 选 x 不选 mat[i][0]\n",
    "                        break  # 剪枝：后面的元素更大，无需枚举\n",
    "                    if dfs(i - 1, s - (x - mat[i][0])):  # 选 x 不选 mat[i][0]\n",
    "                        return True  # 找到 k 个就一直返回 True，不再递归\n",
    "                return False\n",
    "            return dfs(len(mat) - 1, s - sl)  # 先把第一列的所有数都选上\n",
    "\n",
    "        sl = sum(row[0] for row in mat)\n",
    "        sr = sum(row[-1] for row in mat)\n",
    "        return sl + bisect_left(range(sl, sr), True, key=check)\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        a = mat[0][:k]\n",
    "        for row in mat[1:]:\n",
    "            a = sorted(x + y for x in a for y in row)[:k]\n",
    "        return a[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        pre = [0]\n",
    "        for cur in mat:\n",
    "            pre = sorted(a + b for a in pre for b in cur[:k])[:k]\n",
    "        return pre[-1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        m, n = len(mat), len(mat[0])\n",
    "        s = [0]\n",
    "        for i,row in enumerate(mat):\n",
    "            x = sorted(a+b for a,b in itertools.product(s, row))\n",
    "            s = x[:k]\n",
    "        return s[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        print(mat[0][:k])\n",
    "        a= mat[0][:k]\n",
    "        for row in mat[1:]:\n",
    "            print(row)\n",
    "            a = sorted(x + y for x in a for y in row)[:k]\n",
    "            print(a)\n",
    "        return a[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        pre = mat[0][:]\n",
    "        pre.sort()\n",
    "        pre = pre[:k]\n",
    "        for lst in mat[1:]:\n",
    "            cur = []\n",
    "            for num in pre:\n",
    "                for x in lst:\n",
    "                    cur.append(num+x)\n",
    "            cur.sort()\n",
    "            pre = cur[:k]\n",
    "        return pre[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        m = len(mat)\n",
    "        n = len(mat[0])\n",
    "        a = mat[0][:min(k,n)]#第0排前k个数字\n",
    "\n",
    "        for i  in range(1,m):\n",
    "            b = mat[i] #第i行\n",
    "            tmp = []\n",
    "            for x in a:\n",
    "                for y in b:\n",
    "                    tmp.append(x+y)\n",
    "            tmp.sort()\n",
    "            a = tmp[:k]\n",
    "            \n",
    "        return a[-1]\n",
    "\n",
    "    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:\n",
    "        #已经排序，在堆中更新nums2的index\n",
    "        m = len(nums1)\n",
    "        n = len(nums2)\n",
    "        nums= [(nums1[i]+nums2[0],i,0) for i in range(min(k,m))]#已经排序，无需heapify,nums1只需要前个数(k,m)\n",
    "        # print(nums)\n",
    "        # heapify(nums)\n",
    "        res = []\n",
    "        while k and nums:\n",
    "            he,i,j = heappop(nums)\n",
    "            res.append(he)\n",
    "            k -= 1\n",
    "            if j+1<n:\n",
    "                heappush(nums,(nums1[i]+nums2[j+1],i,j+1))\n",
    "            \n",
    "        return res "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        m,n=len(mat),len(mat[0])\n",
    "        def union(a,b):\n",
    "            aa,bb=len(a),len(b)\n",
    "            ls=[]\n",
    "            for i in range(min(aa,k)):\n",
    "                for j in range(min(bb,k)):\n",
    "                    ls.append(a[i]+b[j])\n",
    "            ls.sort()\n",
    "            return ls[:k]\n",
    "        tmp=mat[0]\n",
    "        for i in range(1,m):\n",
    "            tmp=union(mat[i],tmp)\n",
    "        return tmp[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import heapq as hq\n",
    "\n",
    "class Node():\n",
    "    def __init__(self, arr, v):\n",
    "        self.v = v \n",
    "        self.arr = arr\n",
    "    \n",
    "    def sumValue(self):\n",
    "        t = 0\n",
    "        for i in range(len(self.arr)):\n",
    "            a = self.v[i]\n",
    "            t += self.arr[i][a]\n",
    "        return t\n",
    "    \n",
    "    def __gt__(self, other):\n",
    "        return self.sumValue() > other.sumValue()\n",
    "    \n",
    "    def __hash__(self):\n",
    "        r = 0\n",
    "        m = len(self.arr[0]) + 2\n",
    "        for i in range(len(self.arr)):\n",
    "            r = r * m + (self.v[i] + 1)\n",
    "        return r\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        n = len(mat)\n",
    "        m = len(mat[0])\n",
    "        count = 0\n",
    "        queue = [Node(mat, [0] * n)]\n",
    "        \n",
    "        visited = set()\n",
    "        visited.add(queue[0].__hash__())\n",
    "\n",
    "        while queue:\n",
    "            cur = hq.heappop(queue)\n",
    "            count += 1\n",
    "            if count == k:\n",
    "                return cur.sumValue()\n",
    "            \n",
    "            curV = cur.v\n",
    "            for i in range(n):\n",
    "                if curV[i] + 1 < m:\n",
    "                    newV = list(curV)\n",
    "                    newV[i] += 1\n",
    "                    nNode = Node(mat, newV)\n",
    "                    if nNode.__hash__() in visited: continue \n",
    "                    visited.add(nNode.__hash__())\n",
    "                    hq.heappush(queue, nNode)\n",
    "            \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        hq = []\n",
    "        m = len(mat)\n",
    "        n = len(mat[0])\n",
    "        visited = set()\n",
    "        for i in range(min(n,k)):\n",
    "            tmp = [i]\n",
    "            s = mat[0][i]\n",
    "            for j in range(1,m):\n",
    "                s += mat[j][0]\n",
    "                tmp.append(0)\n",
    "            \n",
    "            tmp.insert(0,s)\n",
    "            t = tuple(tmp)\n",
    "            visited.add(t)\n",
    "            hq.append(t)\n",
    "        ans = 0\n",
    "        l = m+1\n",
    "        while k and hq:\n",
    "            tmp = heapq.heappop(hq)\n",
    "            # print(\"选择\",tmp)\n",
    "            ans = tmp[0]\n",
    "            k -= 1\n",
    "            for i in range(2,m+1):\n",
    "                # print(i,newtmp,n+1)\n",
    "                newtmp = list(tmp)\n",
    "                r = i-1\n",
    "                c = newtmp[i]+1\n",
    "                if c<n:\n",
    "                    newtmp[0] -= mat[r][c-1]\n",
    "                    newtmp[0] += mat[r][c]\n",
    "                    newtmp[i] = c\n",
    "                    t = tuple(newtmp)\n",
    "                    if t not in visited:\n",
    "                        visited.add(t)\n",
    "                        heapq.heappush(hq,t)\n",
    "                        # print(newtmp)\n",
    "                    # break\n",
    "        return ans \n",
    "        \n",
    "                \n",
    "\n",
    "            \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        hq = []\n",
    "        m = len(mat)\n",
    "        n = len(mat[0])\n",
    "        visited = set()\n",
    "        for i in range(min(n,k)):\n",
    "            tmp = [i]\n",
    "            s = mat[0][i]\n",
    "            for j in range(1,m):\n",
    "                s += mat[j][0]\n",
    "                tmp.append(0)\n",
    "            tmp.insert(0,s) #将和排在第一位，进行升序排列\n",
    "            t = tuple(tmp)\n",
    "            visited.add(t) #避免重复进堆\n",
    "            hq.append(t)\n",
    "        ans = 0\n",
    "        while k and hq:\n",
    "            tmp = heapq.heappop(hq)\n",
    "            # print(\"选择\",tmp)\n",
    "            ans = tmp[0]\n",
    "            k -= 1\n",
    "            for i in range(2,m+1): #tmp 的长度是m+1，第一个和+行\n",
    "                # print(i,newtmp,n+1)\n",
    "                newtmp = list(tmp)\n",
    "                r = i-1\n",
    "                c = newtmp[i]+1\n",
    "                if c<n:\n",
    "                    newtmp[0] -= mat[r][c-1] #去掉原值\n",
    "                    newtmp[0] += mat[r][c] #加上新值\n",
    "                    newtmp[i] = c #更新位置\n",
    "                    t = tuple(newtmp)\n",
    "                    if t not in visited: #避免重复进堆\n",
    "                        visited.add(t)\n",
    "                        heapq.heappush(hq,t)\n",
    "                        # print(newtmp)\n",
    "                    # break\n",
    "        return ans \n",
    "        \n",
    "                \n",
    "\n",
    "            \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        row=len(mat)\n",
    "        col=len(mat[0])\n",
    "        #40\n",
    "        #k<=200\n",
    "        mid=defaultdict(list)\n",
    "        visit=set()\n",
    "        first=[0]*row\n",
    "        visit.add(\"-\".join([\"0\"]*row))\n",
    "        fsum=sum([x[0] for x in mat])\n",
    "        \n",
    "        cindex=0 \n",
    "        mid[cindex]=first\n",
    "        \n",
    "        hp=[(fsum,cindex)]\n",
    "        small=1\n",
    "        \n",
    "        while small<k:\n",
    "            small+=1\n",
    "            ssum,sindex=heapq.heappop(hp)\n",
    "            preops=mid[sindex]\n",
    "            seops=[str(x) for x in preops]\n",
    "            for i in range(row):\n",
    "                c=int(preops[i])\n",
    "                if c+1<col:\n",
    "                    cops=list(seops)\n",
    "                    cops[i]=str(c+1)\n",
    "                    nkey=\"-\".join(cops)\n",
    "                    if nkey not in visit:\n",
    "                        visit.add(nkey)\n",
    "                        cindex+=1\n",
    "                        nsum=ssum-mat[i][c]+mat[i][c+1]\n",
    "                        mid[cindex]=cops \n",
    "                        heapq.heappush(hp,(nsum,cindex))\n",
    "        # print(hp)\n",
    "        # print(visit)\n",
    "        res,rindex=heapq.heappop(hp)    \n",
    "        return res "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        h= []  # 初始化堆\n",
    "        # mat: m * n  [[1,3,11], [2,4,6]]\n",
    "        # sum(vec[0] for vec in mat):矩阵matrix每一行的首项和\n",
    "        # tuple([0] * len(mat)：初始化一个长度为m的全部为0的数组\n",
    "        # (3, (0, 0))\n",
    "        cur = (sum(vec[0] for vec in mat), tuple([0] * len(mat)))\n",
    "        heapq.heappush(h, cur)  # 入堆\n",
    "        # 为什么cur必须是元组，不能是List[List[int]]? 因为 List[List[int]]无法转化为 集合\n",
    "        seen = set(cur)  # 对 {(0,1),(0,2)}这类的索引串去重\n",
    "        new_sum = 0  # 数组和\n",
    "        \n",
    "        for _ in range(k):  # 共出堆 k 次，刚好是第k小的值\n",
    "            new_sum, pointers = heapq.heappop(h)  # pointers: (2,0)\n",
    "            for i, pointer in enumerate(pointers):\n",
    "                #  Pointer不能等于单个数组的长度-1, 也就是2\n",
    "                # 如果等于2，说明索引到头了，不能再往后移动\n",
    "                if pointer != len(mat[0]) - 1:\n",
    "                    t = list(pointers)  # 把元祖转成list, 是为了修改里面的元素值，因为元组不能修改\n",
    "                    t[i] = pointer + 1  # 把第i个位置的指针, 往后移动1位\n",
    "                    tt = tuple(t)  # 修改完之后，再转回元组\n",
    "                    \n",
    "                    if tt not in seen:  # 避免重复\n",
    "                        seen.add(tt)\n",
    "                        # new_sum 减去当前值，加上下一个值\n",
    "                        heapq.heappush(h, (new_sum - mat[i][pointer] + mat[i][pointer + 1], tt))\n",
    "        return new_sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\"\"\"\n",
    "堆的使用场景\n",
    "三个技巧之一\n",
    "② 多指针（多路归并）解法\n",
    "\"\"\"\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        # 初始化最小堆\n",
    "        h = []\n",
    "        # sum(vec[0] for vec in mat 是m行一维数组的首项和，即mat每行第一个元素的总和\n",
    "        # [0] * len(mat) 初始化一个长度为m的指针，用来指向mat每一行的下标，每一个指针初始化为0，即从下标0开始\n",
    "        # 将上面两项转换成tuple\n",
    "        cur = (sum(vec[0] for vec in mat), tuple([0] * len(mat)))\n",
    "        # 入堆（最小堆）\n",
    "        heapq.heappush(h, cur)\n",
    "        \n",
    "        seen = set(cur)\n",
    "        # 遍历k次，第k次从堆中取出的元素，就是第k小数组和\n",
    "        for _ in range(k):\n",
    "            # total_sum 是当前指针指向的值的数组和\n",
    "            # pointers 是当前每一行指针指向的下标\n",
    "            total_sum, pointers = heapq.heappop(h)\n",
    "            # 粗暴解决: \n",
    "            # 每次都移动指针数组中的一个指针，mat数组总共有n列，(即每行数组长度为n)，故每一个point的可以移动n次，\n",
    "            # 每移动一个指针，就会分叉一次。\n",
    "            for i, pointer in enumerate(pointers):\n",
    "                # pointer不是最后一个位置的指针\n",
    "                if pointer != len(mat[0]) - 1:\n",
    "                    # 将tuple转化为list\n",
    "                    t = list(pointers)\n",
    "                    # 更新指针（右移一位）\n",
    "                    t[i] = pointer + 1\n",
    "                    tt = tuple(t)\n",
    "                    # 若该指针没有出现过\n",
    "                    if tt not in seen:\n",
    "                        seen.add(tt)\n",
    "                        # 更新total_sum， 入堆\n",
    "                        # total_sum加上当前指针指向的值，去掉上一个指针指向的值。\n",
    "                        heapq.heappush(h, (total_sum + mat[i][pointer + 1] - mat[i][pointer], tt))\n",
    "        \n",
    "        return total_sum\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        # 采用小根堆的方法 多路归并\n",
    "        h = [] \n",
    "        element = [sum(value[0] for value in mat),tuple([0]*len(mat))]\n",
    "        heapq.heappush(h,element)\n",
    "        seen = set(element)\n",
    "\n",
    "        for _ in range(k):\n",
    "            acc, pointers = heapq.heappop(h)\n",
    "\n",
    "            for i,pointer in enumerate(pointers):\n",
    "                if pointer != len(mat[0])-1:\n",
    "                    t = list(pointers)\n",
    "                    t[i] = pointer + 1\n",
    "                    tt = tuple(t)\n",
    "                    if tt not in seen:\n",
    "                        seen.add(tt)\n",
    "                        heapq.heappush(h,[acc-mat[i][pointer]+mat[i][pointer+1],tt])\n",
    "        return acc \n",
    "                    \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        m, n = len(mat), len(mat[0])\n",
    "        count = 0\n",
    "        queue = []\n",
    "        heapq.heappush(queue,(sum([mat[i][0] for i in range(m)]), [0] * m))\n",
    "        seen = set()\n",
    "        while count<k:\n",
    "            val, pointers = heapq.heappop(queue)\n",
    "            count+=1\n",
    "            for i in range(m):\n",
    "                if pointers[i]+1<n:\n",
    "                    new_pointers = pointers.copy()\n",
    "                    new_pointers[i] += 1\n",
    "                    if tuple(new_pointers) not in seen:\n",
    "                        heapq.heappush(queue,(val-mat[i][pointers[i]]+mat[i][pointers[i]+1],new_pointers))\n",
    "                        seen.add(tuple(new_pointers))\n",
    "        return val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        matrix = mat\n",
    "        n = len(matrix[0])\n",
    "        m = len(matrix)\n",
    "        useset=set()\n",
    "        useset.add(tuple([0]*m))\n",
    "        rls=[]\n",
    "        v1=sum([matrix[i][0] for i in range(m)])\n",
    "        heap=[(v1,[0]*m)]\n",
    "        for i in range(k ):\n",
    "            if not heap:\n",
    "                return rls[-1]\n",
    "            v,ls=heappop(heap)\n",
    "            gls=[]\n",
    "            rls.append(v)\n",
    "            for i in range(m):\n",
    "                if ls[i]<n-1:\n",
    "                    gls[:]=ls\n",
    "                    gls[i]+=1\n",
    "                    if tuple(gls) not in useset:\n",
    "                        useset.add(tuple(gls))\n",
    "                        xs=[]\n",
    "                        xs[:]=gls\n",
    "                        heappush(heap,(v+matrix[i][gls[i]]-matrix[i][gls[i]-1],xs))\n",
    "\n",
    "        return rls[-1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import heapq\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        hp = []\n",
    "        st = set()\n",
    "        s = 0\n",
    "        m = len(mat)\n",
    "        n = len(mat[0])\n",
    "\n",
    "        for i in range(m):\n",
    "            s += mat[i][0]\n",
    "        \n",
    "        if k == 1:\n",
    "            return s\n",
    "        \n",
    "        for i in range(m):\n",
    "            tmp_0 = [s-mat[i][0]+mat[i][1]]\n",
    "            tmp_1 = [0] * m\n",
    "            tmp_1[i] = 1\n",
    "            st.add(tuple(tmp_1))\n",
    "            heapq.heappush(hp, tmp_0+tmp_1)\n",
    "        \n",
    "        for _ in range(k-1):\n",
    "            tmp = heapq.heappop(hp)\n",
    "\n",
    "            s = tmp[0]\n",
    "            index = tmp[1:]\n",
    "            for i in range(m):\n",
    "                j = index[i]\n",
    "                if j < n - 1:\n",
    "                    tmp_0 = [s - mat[i][j] + mat[i][j+1]]\n",
    "                    tmp_1 = index[:]\n",
    "                    tmp_1[i] += 1\n",
    "                    tp = tuple(tmp_1)\n",
    "                    if tp not in st:\n",
    "                        heapq.heappush(hp, tmp_0+tmp_1)\n",
    "                        st.add(tp)\n",
    "        \n",
    "        return s\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        m,n = len(mat), len(mat[0])\n",
    "        vis = set()\n",
    "        sums = 0\n",
    "        vec = [0] * m \n",
    "        for i in range(m):\n",
    "            sums += mat[i][0]\n",
    "        Q = [(sums, vec)]\n",
    "        cnt = 0\n",
    "        while Q:\n",
    "            sums, vec = heapq.heappop(Q)\n",
    "            cnt += 1\n",
    "            if cnt >= k:\n",
    "                return sums \n",
    "            for i in range(m):\n",
    "                if vec[i] < n-1:\n",
    "                    nvec = copy.deepcopy(vec)\n",
    "                    nvec[i] += 1\n",
    "                    nsums = sums + mat[i][vec[i]+1]-mat[i][vec[i]]\n",
    "                    t = tuple(nvec)\n",
    "                    if t not in vis:\n",
    "                        vis.add(t)\n",
    "                        heapq.heappush(Q, (nsums, nvec))\n",
    "        \n",
    "        return -1 \n",
    "                        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        m, n = len(mat), len(mat[0])\n",
    "        count = 0\n",
    "        queue = []\n",
    "        heapq.heappush(queue,(sum([mat[i][0] for i in range(m)]), [0] * m))\n",
    "        seen = set()\n",
    "        print(n)\n",
    "        while count<k:\n",
    "            print(count,len(queue))\n",
    "            val, pointers = heapq.heappop(queue)\n",
    "            count+=1\n",
    "            for i in range(m):\n",
    "                print(pointers)\n",
    "                if pointers[i]+1<n:\n",
    "                    new_pointers = pointers.copy()\n",
    "                    new_pointers[i] += 1\n",
    "                    if tuple(new_pointers) not in seen:\n",
    "                        heapq.heappush(queue,(val-mat[i][pointers[i]]+mat[i][pointers[i]+1],new_pointers))\n",
    "                        seen.add(tuple(new_pointers))\n",
    "        return val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        h=[]\n",
    "        m,n=len(mat),len(mat[0])\n",
    "        idx=[0]*m  \n",
    "        s=sum(mat[i][0] for i in range(m))\n",
    "        heappush(h,(s,idx))\n",
    "        vis=set()\n",
    "        vis.add(tuple(idx))\n",
    "        cnt=0\n",
    "        while h:\n",
    "            cnt+=1\n",
    "            s,idx=heappop(h)   \n",
    "            if cnt==k:\n",
    "                return s\n",
    "            for i in range(m):\n",
    "                temp=idx[i]\n",
    "                if idx[i]+1<n:\n",
    "                    idx[i]+=1\n",
    "                if tuple(idx) not in vis:\n",
    "                    heappush(h,(s-mat[i][temp]+mat[i][idx[i]],idx.copy()))\n",
    "                    vis.add(tuple(idx))\n",
    "                idx[i]=temp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from heapq import heapify, heappush, heappop\n",
    "from typing import List\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        m, n = len(mat), len(mat[0])\n",
    "        # 用于存储和以及每一行选出的元素的索引的堆\n",
    "        heap = [(sum(row[0] for row in mat), [0]*m)]\n",
    "        heapify(heap)\n",
    "        visited = set(tuple([0]*m))\n",
    "        \n",
    "        while k > 0:\n",
    "            # 弹出当前最小的和\n",
    "            curr_sum, indices = heappop(heap)\n",
    "            k -= 1\n",
    "            \n",
    "            # 如果这是第 k 个最小的和，返回它\n",
    "            if k == 0:\n",
    "                return curr_sum\n",
    "            \n",
    "            # 生成新的可能和\n",
    "            for i in range(m):\n",
    "                new_indices = list(indices)\n",
    "                if new_indices[i] < n - 1:\n",
    "                    new_indices[i] += 1\n",
    "                    new_sum = curr_sum - mat[i][indices[i]] + mat[i][new_indices[i]]\n",
    "                    \n",
    "                    # 避免重复计算\n",
    "                    t_indices = tuple(new_indices)\n",
    "                    if t_indices not in visited:\n",
    "                        heappush(heap, (new_sum, new_indices))\n",
    "                        visited.add(t_indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from heapq import heappush, heappop\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        heap = []\n",
    "        visited = set()\n",
    "        R, C = len(mat), len(mat[0])\n",
    "        res = 0\n",
    "\n",
    "        if not mat:\n",
    "            return -1\n",
    "        idxLst = [0] * R\n",
    "        visited.add(tuple(idxLst))\n",
    "        heappush(heap, [sum([mat[i][0] for i in range(R)]), idxLst])\n",
    "\n",
    "        while heap and k > 0:\n",
    "            res, lst = heappop(heap)\n",
    "            # print(res)\n",
    "            for i, p in enumerate(lst):\n",
    "                if lst[i] < C - 1:\n",
    "                    newLst = lst[:]\n",
    "                    newLst[i] += 1\n",
    "                    newTuple = tuple(newLst)\n",
    "                    if newTuple not in visited:\n",
    "                        newSum = sum([mat[j][newLst[j]] for j in range(R)])\n",
    "                        heappush(heap, [newSum, newLst])\n",
    "                        visited.add(newTuple)\n",
    "            k -= 1\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        m, n = len(mat), len(mat[0])\n",
    "        heap = [(sum([e[0] for e in mat]), [0]*m)]\n",
    "        seen = set([tuple([0]*m)])\n",
    "        while k:\n",
    "            sumnum, now = heappop(heap)\n",
    "            k -= 1\n",
    "            for i in range(m):\n",
    "                if now[i] != n-1:\n",
    "                    new_now = deepcopy(now)\n",
    "                    new_now[i] += 1\n",
    "                    if tuple(new_now) not in seen:\n",
    "                        heappush(heap, (sumnum-mat[i][now[i]]+mat[i][now[i]+1], new_now))\n",
    "                        seen.add(tuple(new_now))\n",
    "        return sumnum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        m, n = len(mat), len(mat[0])\n",
    "        mn = sum(mat[i][0] for i in range(m))\n",
    "        heap = [[mn] + [0] * m]\n",
    "        visited = set(tuple([0] * m))\n",
    "        for _ in range(k):\n",
    "            group = heappop(heap)\n",
    "            for i in range(m):\n",
    "                if group[i + 1] + 1 < n:\n",
    "                    tmp = group.copy()\n",
    "                    delta = mat[i][group[i + 1] + 1] - mat[i][group[i + 1]]\n",
    "                    tmp[0] += delta\n",
    "                    tmp[i + 1] += 1\n",
    "                    if tuple(tmp[1 : ]) not in visited:\n",
    "                        heappush(heap, tmp)\n",
    "                        visited.add(tuple(tmp[1 : ]))\n",
    "        return group[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import heapq\n",
    "\n",
    "class Solution:\n",
    "    def kthSmallest(self, mat: List[List[int]], k: int) -> int:\n",
    "        n = len(mat)\n",
    "        m = len(mat[0])\n",
    "        indications = []\n",
    "        base = 0\n",
    "        for i in range(n):\n",
    "            base += mat[i][0]\n",
    "        idx = tuple(0 for _ in range(n))\n",
    "        heapq.heappush(indications, (base, *idx))\n",
    "        rst = base\n",
    "        visited = set()\n",
    "        visited.add(idx)\n",
    "        while k > 0:\n",
    "            min_rst = heapq.heappop(indications)\n",
    "            rst = min_rst[0]\n",
    "            for i in range(n):\n",
    "                if min_rst[i + 1] + 1 < m:\n",
    "                    idx = tuple(min_rst[x + 1] + (0 if x != i else 1) for x in range(n))\n",
    "                    if idx not in visited:\n",
    "                        visited.add(idx)\n",
    "                        heapq.heappush(indications, (min_rst[0] - mat[i][min_rst[i + 1]] + mat[i][min_rst[i + 1] + 1], *idx))\n",
    "            k -= 1\n",
    "        return rst"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
